import camelot
import fitz  # PyMuPDF
from PIL import Image, ImageDraw
from pathlib import Path
from PyPDF2 import PdfReader  # 替换 PdfFileReader

def extract_table_coordinates(pdf_path):
    # 读取PDF文件
    tables = camelot.read_pdf(pdf_path, pages='all')

    table_coords = {}
    for table in tables:
        page_num = table.page
        if page_num not in table_coords:
            table_coords[page_num] = []
        table_coords[page_num].append((table.x1, table.y1, table.x2, table.y2))
    return table_coords

def convert_pdf_to_images(pdf_path):
    pdf_document = fitz.open(pdf_path)
    images = []
    for page_num in range(len(pdf_document)):
        page = pdf_document.load_page(page_num)  # 加载页面
        pix = page.get_pixmap()  # 将页面转换为图像
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        images.append(img)
    return images

def draw_table_boxes_on_images(images, table_coords, save:Path):
    for page_num, img in enumerate(images):
        draw = ImageDraw.Draw(img)
        if (page_num + 1) in table_coords:
            for box in table_coords[page_num + 1]:
                x1, y1, x2, y2 = box
                draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
        img.save(save.joinpath(f'page_{page_num + 1}_with_tables.png'))

def main(pdf_path):
    save = Path(pdf_path).parent.joinpath(Path(pdf_path).stem)
    save.mkdir(parents=True, exist_ok=True)
    table_coords = extract_table_coordinates(pdf_path)
    images = convert_pdf_to_images(pdf_path)
    draw_table_boxes_on_images(images, table_coords, save)

pdf_path = 'Company/云鹰_MEC边缘计算单元_竖版彩页.pdf'
main(pdf_path)